package org.ysh.core.ocr;

import com.benjaminwan.ocrlibrary.OcrResult;
import org.apache.commons.io.FilenameUtils;
import org.ysh.annotation.OcrDetector;
import org.ysh.core.utils.ClassUtils;
import org.ysh.core.utils.StringUtils;

import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CompletableFuture;

/**
 * OCR工厂类
 */
public class OcrDetectFactory {

    private static final Map<String,OcrDetect> DETECT_MAP = new HashMap<>();

    public OcrDetectFactory(String ... resolverPackages ){
        String defaultPackage = "com.ruoyi.common.core.utils.ocr";
        boolean contains = false;
        if (StringUtils.isNotEmpty(resolverPackages)) {
            for (String pack : resolverPackages) {
                if (defaultPackage.equals(pack)) {
                    contains = true;
                    break;
                }
            }

            for (String packageName : resolverPackages) {
                Set<Class<?>> classSet = ClassUtils.getClasses(packageName);
                loadImageOcrDetect(classSet);
            }
        }

        if (!contains) {
            Set<Class<?>> classSet = ClassUtils.getClasses(defaultPackage);
            loadImageOcrDetect(classSet);
        }
    }

    private void loadImageOcrDetect(Set<Class<?>> classes) {
        Class<OcrDetect> detectorClass = OcrDetect.class;
        if (null != classes && !classes.isEmpty()) {
            for (Class<?> clazz : classes) {
                if (Arrays.stream(clazz.getInterfaces()).anyMatch(interfaceClass->interfaceClass == detectorClass)) {
                    OcrDetector detector = clazz.getDeclaredAnnotation(OcrDetector.class);
                    if (detector != null) {
                        String[] extensions = detector.fileExt();
                        try {
                            OcrDetect info = (OcrDetect) clazz.getConstructor().newInstance();
                            if (StringUtils.isNotEmpty(extensions)) {
                                for (String ext : extensions) {
                                    DETECT_MAP.put(ext, info);
                                }
                            }
                        } catch (Exception e) {
                            e.printStackTrace();
                        }
                    }
                }
            }
        }
    }

    /**
     * 识别附件
     * @param filePath
     * @return
     * @throws Exception
     */
    public OcrResult detectFile(String filePath) throws Exception {
        String extension = FilenameUtils.getExtension(filePath);
        OcrDetect detect = DETECT_MAP.get(extension);
        if(detect == null){
            return null;
        }
        return CompletableFuture.supplyAsync(() -> detect.detect(filePath)).get();
    }

    public static void main(String[] args) throws Exception {
        OcrDetectFactory factory = new OcrDetectFactory();
        OcrResult ocrResult = factory.detectFile("D:\\BaiduNetdiskDownload\\5179\\001.jpg");
        System.out.println(ocrResult);
    }

}
